import os
import streamlit as st
from large_model.agent_executors import Executors
from streamlit_model.streamlit_utils.file_uploader import read_uploaded_file
from streamlit_model.streamlit_utils.session_utils import load_session, save_session

PAGE_KEY = "uploaded_file_case"

def uploaded_file_case():
    st.write("根据需求文档按照要求生成用例")

    cached = load_session(PAGE_KEY)

    if st.session_state.pop("clear_triggered", False):
        st.session_state.clear()
        st.rerun()

    if "task_executed" not in st.session_state:
        st.session_state.update(cached)

    if st.session_state.get("task_executed"):
        _, center_output, _ = st.columns([1, 3, 1])
        with center_output:
            st.success("测试用例生成完成！")
            st.markdown(st.session_state["output"])

    cached_file_name   = cached.get("uploaded_file_name")
    cached_review      = cached.get("review_method", "需求文档是单个模块完整的功能")
    cached_desc        = cached.get("description", "")
    cached_content     = cached.get("uploaded_file_content", "")
    cached_file_path   = cached.get("file_path")

    _, upload_col, _ = st.columns([1, 3, 1])
    with upload_col:
        st.write("当前的能力只能处理小文档，如果是大文档需要优化（分块处理）")
        uploaded_file = st.file_uploader("选择需求文档进行上传，只能上传一个文件",
                                         type=["doc", "docx", "pdf"])

    # 读取文件内容
    file_content = ""
    file_name = None
    file_path = None
    if uploaded_file is not None:
        file_bytes = uploaded_file.read()
        file_name = uploaded_file.name
        file_content = read_uploaded_file(file_bytes, file_name)
    elif cached_file_name and cached_content:
        file_content = cached_content
        file_name = cached_file_name
        file_path = cached_file_path

    _, sel_col, _ = st.columns([1, 3, 1])
    with sel_col:
        opts = ["需求文档是单个模块完整的功能", "需求文档是连续的", "大文件处理"]
        review_method = st.selectbox("请选择文档处理方式", opts,
                                     index=opts.index(cached_review) if cached_review in opts else 0)

        description = ""
        if review_method == "需求文档是连续的":
            description = st.text_area("请输入描述", value=cached_desc, height=300)
        elif review_method == "大文件处理":
            if uploaded_file is not None:
                temp_dir = "temp_uploads"
                os.makedirs(temp_dir, exist_ok=True)
                file_path = os.path.join(temp_dir, file_name)
                with open(file_path, "wb") as f:
                    f.write(file_bytes)
            else:
                file_path = cached_file_path
            description = st.text_area("输入相关的描述", value=cached_desc, height=300)

    # 即时保存
    need_save = (
        file_name != cached_file_name or
        review_method != cached_review or
        description != cached_desc or
        file_content != cached_content or
        file_path != cached_file_path
    )
    if need_save and not st.session_state.get("clear_triggered", False):
        save_session(PAGE_KEY, {
            "uploaded_file_name": file_name,
            "uploaded_file_content": file_content,
            "review_method": review_method,
            "description": description,
            "file_path": file_path,
            "task_executed": st.session_state.get("task_executed", False),
            "output": st.session_state.get("output", "")
        })

    _, btn_col, _ = st.columns([2, 1, 2])
    with btn_col:
        generate_button = st.button("生成用例", use_container_width=True)
        if st.button("清除缓存", use_container_width=True):
            save_session(PAGE_KEY, {})
            st.session_state["clear_triggered"] = True
            st.rerun()

        if generate_button and file_content:
            executors = Executors()
            with st.spinner("正在生成测试用例..."):
                if review_method == "需求文档是单个模块完整的功能":
                    output = executors.function_mians(file_content)
                elif review_method == "大文件处理":
                    prompt = f"doc_path文件路径：{file_path}，用户要求：{description}"
                    output = executors.function_mian_descriptions(prompt, temp=True)
                else:
                    output = executors.function_mian_description(file_content, description)

            data = {
                "uploaded_file_name": file_name,
                "uploaded_file_content": file_content,
                "review_method": review_method,
                "description": description,
                "file_path": file_path,
                "task_executed": True,
                "output": output
            }
            st.session_state.update(data)
            if not st.session_state.get("clear_triggered", False):
                save_session(PAGE_KEY, data)
            st.rerun()